package com.anonymous.util.test;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import com.anonymous.message.LoginMsg;
import com.anonymous.message.Message;
import com.anonymous.util.DatabaseUtils;
/*
 * Author : Sungkyu Kang
 * Description
 * 		This class is to demonstrate and test security code before applying to the system.
 */
public class SecurityProtocols implements Runnable {
	
	/* 
	 * Related data structures
	 */
	
	/*
	 * To generate RSA keys and save them in files. This method may be executed once, or never be used if keys exsit.
	 * - Mechanisms
	 * Key type : RSA
	 * Padding(Encoding) : PKCS8 (not to have to have separate modulo and exponent)
	 * File names : serverPrivateKey.key, serverPublicKey.key
	 */
	public void generateServerKeys(){
		try {
			KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
			keyPairGen.initialize(3072, SecureRandom.getInstance("SHA1PRNG", "SUN"));
			KeyPair pair = keyPairGen.generateKeyPair();
			byte[] encodedPrivate = pair.getPrivate().getEncoded();
			FileOutputStream fos = new FileOutputStream(new File("serverPrivateKey.key"));
			fos.write(encodedPrivate);
			
			byte[] encodedPublic = pair.getPublic().getEncoded();
			fos = new FileOutputStream(new File("serverPublicKey.key"));
			fos.write(encodedPublic);
			
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchProviderException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (FileNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
	}
	
	/*
	 * To read server's private and public key from files.
	 * File names, algorithm, and specifications are hardcoded so have to change as other mechanisms are chosen.
	 * TODO Allow to choose the files by taking input from console.
	 */
	public void initializeServerCriptoSchem(){
		File serverPublicKeyFile = new File("serverPublicKey.key");
		File serverPrivateKeyFile = new File("serverPrivateKey.key");
		byte[] publicKeyBuf = new byte[(int)serverPublicKeyFile.length()];
		byte[] privateKeyBuf = new byte[(int)serverPrivateKeyFile.length()];
		FileInputStream fis;
		PrivateKey serverPrivateKey;
		PublicKey serverPublicKey;

		try {
			fis = new FileInputStream(serverPublicKeyFile);
			fis.read(publicKeyBuf);

			fis = new FileInputStream(serverPrivateKeyFile);
			fis.read(privateKeyBuf);
			KeyFactory keyFac = KeyFactory.getInstance("RSA");
			
			KeySpec publicKeySpec = new X509EncodedKeySpec(publicKeyBuf);
			serverPublicKey = keyFac.generatePublic(publicKeySpec);
			ServerCriptoSchem.setServerPublicKey(serverPublicKey);

			KeySpec privateKeySpec = new PKCS8EncodedKeySpec(privateKeyBuf);
			serverPrivateKey = keyFac.generatePrivate(privateKeySpec);
			ServerCriptoSchem.setServerPriavateKey(serverPrivateKey);

			
		} catch (FileNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvalidKeySpecException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		
	}
	
	/*
	 * Client method that loads server public key from built-in file.
	 * If fails, ask the server for one.
	 */
	public void loadServerPublicKey(){
		File serverPublickKeyFile = new File("serverPublicKey.key");
		FileInputStream fin;
		try {
			fin = new FileInputStream(serverPublickKeyFile);
		byte[] keyFileBuf = new byte[(int)serverPublickKeyFile.length()];

		fin.read(keyFileBuf);
		KeyFactory keyFac = KeyFactory.getInstance("RSA");
		KeySpec keySpec = new PKCS8EncodedKeySpec(keyFileBuf);
		
		ClientCriptoSchem.setServerPublicKey(keyFac.generatePublic(keySpec));
		} catch (FileNotFoundException e) {
			System.err.println("Server Public key does not exist. Downlaod from server");
			e.printStackTrace();
		} catch (IOException e) {
			System.err.println("IO exception");
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			System.err.println("Alogithm exception");
			e.printStackTrace();
		} catch (InvalidKeySpecException e) {
			System.err.println("Key Spec Exception");
			e.printStackTrace();
		}
	}

	/*
	 * ******************************Login operation cryptographic protocol******************************
	 * 
	 * - Public key of the server is distributed in the client software, and do not need to be updated.
	 * 
	 * 1. Client : SessionKey =Gen(l)
	 * 			ClientPublicKey, ClientPrivate Key = Gen(l)
	 *			c = Enc(id, password, SessionKey ;ServerPublicKey)
	 * 
	 * 2. Client -> Server : c
	 * 
	 * 3. Server : id, password, SessionKey = Dec(c; ServerPrivateKey) 
	 * 			if( succeed authenticating database)
	 * 				//TempSessionKey = Gen(l); for now only client generate session key 
	 * 				//SessionKey = MAC(TempSessionKey | SessionKey)
	 * 				Create user entry in the list, save SessionKey, PublicKey of the user.
	 * 			else toss 
	 * 
	 * 4. Server->Client : Enc(userlist; ClientPublicKey), 
	 * 
	 * 5.Client : userlist = Dec(Enc(userlist; ClientPublicKey), ClientPrivateKey)
	 * 
	 * 
	 * note :  man in the middle attack possible with the given session key?
	 * note : Mechanism part of the implementation must be reconsidered, so the implementation may change accordingly.
	 */
	public void clientLogin(String id, String password){
		
		try {
			int sessionKey;
			KeyPair pair = null;
			PrivateKey myPrivateKey;
			PublicKey myPublicKey;
			PublicKey serverPublicKey = ClientCriptoSchem.getServerPublicKey();
			
			// Generate session key to talk to server; Mechanism- SecureRandom
			sessionKey = SecureRandom.getInstance("SHA1PRNG", "SUN").nextInt();
			// Generate client's RSA keys
			KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
			keyGen.initialize(3072, SecureRandom.getInstance("SHA1PRNG", "SUN"));
			pair = keyGen.generateKeyPair();
			myPrivateKey = pair.getPrivate();
			myPublicKey = pair.getPublic();
			ClientCriptoSchem.setMyPrivateKey(myPrivateKey);			
			//Encrypt id, password, public key, session key
			
			//c = Enc(id, password, SessionKey ;ServerPublicKey)
			Cipher c = Cipher.getInstance("RSA/EBC/PKCS5Padding");
			c.init(Cipher.ENCRYPT_MODE, serverPublicKey);
			
			LoginMsg loginMsg = new LoginMsg();
			//Antony delete something here
			loginMsg.setPasswordHash(password);
			loginMsg.setUsername(id);
			
			/******************************************************************************************************
			 * message encryption and network transfer 
			 * 
			 *  Create Message object
			 *  Convert into a byte array
			 *  Get a cipher(byte array)
			 *  Transfer the cipher through network(byte array)
			 *  Decrypt the cipher(byte array)
			 *  Type convert to Message object
			 *  Evaluate flag, convert into specific type depending on flag
			 * 
			 ******************************************************************************************************/
			// Client -> Server : c
			ByteArrayOutputStream bos = new ByteArrayOutputStream();
			ObjectOutputStream ost = new ObjectOutputStream(bos);
			ost.writeObject(loginMsg);
			byte[] msg = c.doFinal(bos.toByteArray());
			
			Socket soc =  new Socket("12.111.111.111",  4444);
			BufferedOutputStream bufferedOutputStream = new BufferedOutputStream(soc.getOutputStream());
			bufferedOutputStream.write(msg);
			bufferedOutputStream.close();
			
			// skip through server part 
			

			// 5.Client : userlist = Dec(Enc(userlist; ClientPublicKey), ClientPrivateKey)
			byte[] readBuffer = new byte[100];
			BufferedInputStream bis = new BufferedInputStream(soc.getInputStream());
			bis.read(readBuffer);

			ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(readBuffer));
			//clientUserList  = (UserList)ois.readObject();
			
			ois.close();
			ost.close();
			bos.close();
			soc.close();
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}
	
	@Override
	public void run() {
		ServerSocket ss;
		Socket soc;
		byte[] inputBuffer = new byte[100];
		try{

			ss = new ServerSocket(4444);
			while(true){
				soc = ss.accept();
			
			BufferedInputStream bis = new BufferedInputStream(soc.getInputStream());
			bis.read(inputBuffer);
			Cipher c = Cipher.getInstance("RSA/EBC/Nopadding");
			c.init(Cipher.DECRYPT_MODE, ServerCriptoSchem.getServerPriavateKey());
			c.doFinal(inputBuffer);
			
			ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(inputBuffer));
			
			Message message = (Message)ois.readObject();
			//if("LOGIN".equals(message.getFlag()))
			message = (LoginMsg)message;
			

			BufferedOutputStream bus = new BufferedOutputStream(soc.getOutputStream());

			// Authenticate against database
			// This part of the method will be implemented in separate method.
			if(DatabaseUtils.getInstance().login(((LoginMsg)message).getUsername(), ((LoginMsg)message).getPasswordHash()))
			{
				//Save generate random number to create session key, save the session key and client's public key in the list with name, password
				//new User();
				
				
				new String();
				
			}
			else
				new LoginMsg();
			
				
//			c.init(Cipher.ENCRYPT_MODE, /*client's public key*/);


//			bus.write(encryptedBuffer);

			soc.close();
			}
		}catch(IOException e){
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
	}

	/*
	 * This part of the server will be working as a thread accepting client socket message passing. 
	 */
	
	public void serverAcceptLogin(){
		ServerSocket ss;
		Socket soc;
		byte[] inputBuffer = new byte[100];
		try{

			ss = new ServerSocket(4444);
			soc = ss.accept();
			
			BufferedInputStream bis = new BufferedInputStream(soc.getInputStream());
			bis.read(inputBuffer);
			Cipher c = Cipher.getInstance("RSA/EBC/Nopadding");
			c.init(Cipher.DECRYPT_MODE, ServerCriptoSchem.getServerPriavateKey());
			c.doFinal(inputBuffer);
			

			ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(inputBuffer));
			
			Message message = (Message)ois.readObject();
			//if("LOGIN".equals(message.getFlag()))
			message = (LoginMsg)message;
			

			BufferedOutputStream bus = new BufferedOutputStream(soc.getOutputStream());

			// Authenticate against database
			// This part of the method will be implemented in separate method.
			if(DatabaseUtils.getInstance().login(((LoginMsg)message).getUsername(), ((LoginMsg)message).getPasswordHash()))
			{
				//Save generate random number to create session key, save the session key and client's public key in the list with name, password
				//new User();
				
				
				new String();
				
			}
			else
				new LoginMsg();
			
				
//			c.init(Cipher.ENCRYPT_MODE, /*client's public key*/);
			
//			bus.write(encryptedBuffer);

			soc.close();
			
		}catch(IOException e){
			e.printStackTrace();
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
	}
	
	
	public static void main(String[] args) {
		//new SecurityProtocols().generateServerKeys();

		new SecurityProtocols().initializeServerCriptoSchem();
		System.out.println(ServerCriptoSchem.getServerPriavateKey());
		System.out.println(ServerCriptoSchem.getServerPublicKey());

	}

}

/*
 * What a client needs :
 *  Server's RSA public key, its own session key, user list that contains other users public keys session keys.
 *  (Other information such as username is left out here) 
 *  
 */
class  ClientCriptoSchem{
	private static PublicKey serverPublicKey = null;
	private static PrivateKey myPrivateKey = null;
	private static int sessionKey = 0;
	
	public static PublicKey getServerPublicKey() {
		return serverPublicKey;
	}

	public static PrivateKey getMyPrivateKey() {
		return myPrivateKey;
	}

	public static int getSessionKey() {
		return sessionKey;
	}

	public static void setServerPublicKey(PublicKey serverPublicKey) {
		ClientCriptoSchem.serverPublicKey = serverPublicKey;
	}

	public static void setMyPrivateKey(PrivateKey myPrivateKey) {
		ClientCriptoSchem.myPrivateKey = myPrivateKey;
	}

	public static void setSessionKey(int sessionKey) {
		ClientCriptoSchem.sessionKey = sessionKey;
	}

	private ClientCriptoSchem(){
		

	}
	
	
}

/*
* What server needs :
* Server's RSA private key, User list that contains each users public keys and session keys.
 */

class ServerCriptoSchem{
	private static PrivateKey serverPriavateKey;
	private static PublicKey serverPublicKey;

	public static PrivateKey getServerPriavateKey() {
		return serverPriavateKey;
	}
	public static PublicKey getServerPublicKey() {
		return serverPublicKey;
	}
	public static void setServerPriavateKey(PrivateKey serverPriavateKey) {
		ServerCriptoSchem.serverPriavateKey = serverPriavateKey;
	}
	public static void setServerPublicKey(PublicKey serverPublicKey) {
		ServerCriptoSchem.serverPublicKey = serverPublicKey;
	}
	
	private ServerCriptoSchem(){
		
	}

}